# 导入库
import pandas as pd
import statsmodels.formula.api as smf 
from sqlalchemy import create_engine
import pymysql
from matplotlib import pyplot as plt
import seaborn as sns

db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8'
}

engine = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}")
conn = pymysql.connect(**db_config)
chunk_size = 10000

df1 = pd.read_sql_query("""
    SELECT d.*, m.buy_lg_vol, m.sell_lg_vol,m.buy_elg_vol,m.sell_elg_vol,m.net_mf_vol, i.vol as i_vol, i.closes as i_closes from data_1 d
    join moneyflow m on d.ts_code = m.ts_code and d.trade_date = m.trade_date
    left join index_daily i on d.trade_date = i.trade_date and i.ts_code = '399001.SZ'
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' and d.ts_code = '002229.SZ'
    
""", conn,chunksize=chunk_size)

df1 = pd.concat(df1, ignore_index=True)
print(df1.head)

df1['zd_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
df1['i_closes'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1), 2)
df1['i_vol'] = round((df1['i_vol'] - df1['i_vol'].shift(1)) / df1['i_vol'].shift(1), 2)

# 处理异常
df1 = df1.dropna(subset=['rd_closes', 'rz_closes', 'rz_vol'])
ax = df1.plot(x='trade_date', y=['closes', 'rz_closes', 'rz_vol', 'high', 'low', 'closes', 'rz_closes', 'changes', 'net_mf'])

df1 = df1.dropna(include=['number']).columns.to_list()
number_list = [i for i in number if i not in m]

formula = 'rd_closes ~ ' + ' + '.join(number_list)
res = smf.ols(formula, data=df1).fit()
print(res.summary())

plt.figure(figsize=(15, 8))
plt.scatter(df1['rz_closes'], df1['closes'], data=df1)
plt.show()

features = df1[['rz_vol', 'amount', 'buy_lwol', 'sell_lwol', 'buy_mall', 'sell_mall', 'net_mf_vol', 'zs_vol']]
correlation_matrix = features.corr()
print(correlation_matrix)

plt.figure(figsize=(10, 8))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', fmt='.2f', linewidths=5)
plt.title('Correlation Matrix')
plt.show()
